import json
import requests
import traceback
from typing import List, Dict, Any
from db_manager import DatabaseManager
from ai_config import AI_CONFIG, PROMPT_TEMPLATE

class AIQueryManager:
    def __init__(self):
        self.db = DatabaseManager()
        self.table_info = self._load_table_info()
        self.headers = {
            "Authorization": f"Bearer {AI_CONFIG['api_key']}",
            "Content-Type": "application/json"
        }
        # print("self.table_info",self.table_info)
    
    def _clean_sql(self, sql: str) -> str:
        """
        清理SQL语句，移除markdown标记等
        """
        # 移除可能的markdown代码块标记
        sql = sql.replace('```sql', '').replace('```', '').strip()
        return sql
    
    def _get_column_names(self, table_name: str) -> List[str]:
        """
        获取表的列名列表
        """
        if table_name in self.table_info:
            return [col['column_name'] for col in self.table_info[table_name]]
        return []
    
    def _validate_and_fix_sql(self, sql: str) -> str:
        """
        验证并修复SQL语句
        """
        sql = self._clean_sql(sql)
        
        # 修复常见的列名错误
        sql = sql.replace('status = 1', 'is_deleted = 0')  # 状态判断改为is_deleted判断
        sql = sql.replace('tenant_id', 'id')  # 修复租户ID的引用
        
        # 移除不存在的字段条件
        if 'system_user' in sql and 'tenant_id' in sql:
            # 如果是查询用户和租户的关联，使用system_user_role表
            sql = sql.replace(
                'JOIN system_tenant st ON su.tenant_id = st.id',
                'JOIN system_user_role sur ON su.id = sur.user_id JOIN system_tenant st ON sur.tenant_id = st.id'
            )
        
        return sql
    
    def _load_table_info(self) -> dict:
        """
        加载数据库表结构信息
        """
        tables = self.db.get_tables()
        table_info = {}
        
        for table in tables:
            table_name = table['table_name']
            table_comment = table['table_comment'] or f"表 {table_name}"
            columns = self.db.get_table_columns(table_name)
            
            # 构建字段描述列表
            column_descriptions = []
            for col in columns:
                col_name = col['column_name']
                col_type = col['data_type']
                col_nullable = "可空" if col['is_nullable'] == 'YES' else "非空"
                col_comment = col['column_comment'] or f"字段 {col_name}"
                column_descriptions.append(f"{col_name} ({col_type}, {col_nullable}): {col_comment}")
            
            # 构建表描述
            table_info[table_name] = {
                'description': table_comment,
                'columns': column_descriptions
            }
        
        return table_info
    
    def _generate_table_description(self) -> str:
        """
        生成数据库表结构的描述文本
        """
        table_info = self._load_table_info()
        description = "数据库包含以下表：\n\n"
        
        for table_name, info in table_info.items():
            description += f"表名：{table_name}\n"
            description += f"描述：{info['description']}\n"
            description += "字段列表：\n"
            for col_desc in info['columns']:
                description += f"  - {col_desc}\n"
            description += "\n"
        
        return description
    
    def _generate_query_prompt(self, question: str) -> str:
        """
        生成查询提示
        """
        table_description = self._generate_table_description()
        return PROMPT_TEMPLATE.format(
            table_description=table_description,
            question=question
        )
    
    def _get_ai_response(self, prompt: str) -> str:
        """
        调用DeepSeek API获取响应
        """
        try:
            print(f"正在调用DeepSeek API...")
            response = requests.post(
                f"{AI_CONFIG['api_base']}/chat/completions",
                headers=self.headers,
                json={
                    "model": AI_CONFIG['model'],
                    "messages": [
                        {"role": "system", "content": "你是一个SQL专家，负责将自然语言转换为SQL查询语句。请只返回SQL语句，不要包含任何markdown标记或其他说明。\n"
                         +"请注意针对PostgreSQL数据库，请使用正确的语法。例如涉及到user等是PostgreSQL的保留关键字相关的数据库表名或字段，需要使用双引号包裹。（例如虽然是system_user，但 PostgreSQL 在某些解析情况下仍然容易对类似的名字卡壳，需要使用双引号包裹system_user）."
                         +"请注意正确理解个字段的数据类型，避免错误数据转换。\n"},
                        {"role": "user", "content": prompt}
                    ],
                    "temperature": AI_CONFIG['temperature'],
                    "max_tokens": AI_CONFIG['max_tokens']
                }
            )
            
            if response.status_code == 200:
                result = response.json()
                return result['choices'][0]['message']['content'].strip()
            else:
                error_msg = f"API调用失败: {response.status_code} - {response.text}"
                print(error_msg)
                return None
                
        except Exception as e:
            error_msg = f"AI模型调用失败: {str(e)}"
            print(error_msg)
            print(traceback.format_exc())
            return None
    
    def query(self, question: str) -> Dict[str, Any]:
        """
        处理自然语言查询
        """
        try:
            # 生成查询提示
            prompt = self._generate_query_prompt(question)
            
            # 调用AI模型获取SQL查询
            sql_query = self._get_ai_response(prompt)
            
            if not sql_query:
                return {
                    'success': False,
                    'error': '无法生成SQL查询',
                    'traceback': traceback.format_exc()
                }
            
            # 验证和修复SQL
            sql_query = self._validate_and_fix_sql(sql_query)
            print(f"处理后的SQL: {sql_query}")
            
            # 执行SQL查询
            results = self.db.execute_query(sql_query)
            return {
                'success': True,
                'sql': sql_query,
                'results': results
            }
        except Exception as e:
            return {
                'success': False,
                'error': str(e),
                'traceback': traceback.format_exc()
            }
    
    def get_table_info(self) -> Dict[str, List[Dict[str, str]]]:
        """
        获取表结构信息
        """
        return self.table_info 